'''
Created on Nov 11, 2009

@author: mkiyer
'''

import numpy as np

def get_joint_dist(xarr, yarr, x_binsize, y_binsize):
    '''
    given numpy arrays xarr and yarr with respective bin sizes x_binsize
    and y_binsize, compute the empirical pdf f(x,y)
    
    returns an empirical distribution matrix along with bin sizes needed
    to compute individual probabilities
    '''
    # function to generate equal sized bins over the range of an array's values
    abin = lambda a, bin_size: np.arange(0, bin_size * (1 + np.ceil(np.max(np.abs(a))/float(bin_size))), bin_size)
    mybins = (abin(xarr, x_binsize),
              abin(yarr, y_binsize))
    h, xedges, yedges = np.histogram2d(np.abs(xarr), np.abs(yarr), bins=mybins, normed=True)
    return h, x_binsize, y_binsize    
    
def get_p_value(empirical_dist, x, y):
    '''
    given an empirical distribution f(x,y) with bin sizes and 
    the values x and y, return the probability of P(X >= x, Y >= y)
    
    -the 'empirical_dist' parameter is the return value of the 
     get_joint_dist() function.  it is a 3-tuple (h, x_binsize, y_binsize).
    '''
    h, x_binsize, y_binsize = empirical_dist
    xbin = int(float(np.abs(x) / x_binsize))
    ybin = int(float(np.abs(y) / y_binsize))
    return (1.0 - (np.sum(h[0:xbin,0:ybin]) * x_binsize * y_binsize))

if __name__ == '__main__':
    xarr = np.array([0, 0, 0, 0, 1, 1, 2, 2, 5, 5, 5, 1, 1, 10])
    yarr = np.array([1, 1, 1, 2, 3, 4, 5, 5, 5, 6, 7, 10, 20, 30])
    empdist = get_joint_dist(xarr, yarr, 1, 1)
    print get_p_value(empdist, 9, 20)
    